"""
cspdarknet: backbone for yolov4 and yolov5
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# Darknet and cspdarknet are relatively self-contained, do not use other layers, 
# define layers here.

from vortex.network.layers.common import get_activation, BasicConv


class ResBlock(nn.Module):
    """
    cspdarknet use Mish activation, do not provide choices.
    """
    def __init__(self, channels, hidden_channels=None, act='mish'):
        super(ResBlock, self).__init__()
        if hidden_channels is None:
            hidden_channels = channels
        self.conv1 = BasicConv(channels, hidden_channels, 1, act='mish')
        self.conv2 = BasicConv(hidden_channels, channels, 3, act=None)
        self.activation = get_activation(act)
    
    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = x + out
        out = self.activation(out)
        return out


class CSPFirst(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(CSPFirst, self).__init__()
        self.downsample = BasicConv(in_channels, out_channels, 3, stride=2, act='mish')
        self.trans0 = BasicConv(out_channels, out_channels, 1, act='mish')
        self.trans1 = BasicConv(out_channels, out_channels, 1, act='mish')
        # self.blocks = nn.Sequential(
        #     ResBlock(channels=out_channels, hidden_channels=out_channels//2),
        #     BasicConv(out_channels, out_channels, 1)
        # )
        self.blocks = ResBlock(channels=out_channels, hidden_channels=out_channels//2)
        
        self.trans_cat = BasicConv(out_channels*2, out_channels, 1, act='mish')
    
    def forward(self, x):
        x = self.downsample(x)
        x0 = self.trans0(x)
        x1 = self.trans1(x)
        x1 = self.blocks(x1)
        x = torch.cat([x0, x1], dim=1)
        x = self.trans_cat(x)
        return x


class CSPStem(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks):
        super(CSPStem, self).__init__()
        self.downsample = BasicConv(in_channels, out_channels, 3, stride=2, act='mish')
        self.trans0 = BasicConv(out_channels, out_channels//2, 1, act='mish')
        self.trans1 = BasicConv(out_channels, out_channels//2, 1, act='mish')
        self.blocks = nn.Sequential(
            *[ResBlock(out_channels//2) for _ in range(num_blocks)]
        )
        self.trans_cat = BasicConv(out_channels, out_channels, 1, act='mish')

    def forward(self, x):
        x = self.downsample(x)
        x0 = self.trans0(x)
        x1 = self.trans1(x)
        x1 = self.blocks(x1)
        x = torch.cat([x0, x1], dim=1)
        x = self.trans_cat(x)
        return x


class CSPDarknet(nn.Module):
    def __init__(self, layers):
        super(CSPDarknet, self).__init__()
        self.channels = [64, 128, 256, 512, 1024]
        self.inplanes = 32
        self.conv1 = BasicConv(3, self.inplanes, kernel_size=3, stride=1)
        self.stage0 = CSPFirst(32, self.channels[0])
        self.stage1 = CSPStem(self.channels[0], self.channels[1], layers[0])
        self.stage2 = CSPStem(self.channels[1], self.channels[2], layers[1])
        self.stage3 = CSPStem(self.channels[2], self.channels[3], layers[2])
        self.stage4 = CSPStem(self.channels[3], self.channels[4], layers[3])

        self.init_weights()
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.stage0(x)
        x = self.stage1(x)
        out1 = self.stage2(x)
        out2 = self.stage3(out1)
        out3 = self.stage4(out2)
        return out1, out2, out3

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


def cspdarknet53(pretrained=None):
    model = CSPDarknet([2, 8, 8, 4])
    if pretrained is not None:
        state_dict = torch.load(pretrained, map_location='cpu')
        model.load_state_dict(state_dict)
    return model
